"""
Author: '夜微凉'
Date: 2025-01-09 10:53:53
LastEditors: '夜微凉'
LastEditTime: 2025-01-09 11:04:26
FilePath: /lsq-net-master/lrg/utils/data_loader.py
Description: 

"""

import os

import numpy as np
import torch
import torch.utils.data
import torchvision
from sklearn.model_selection import train_test_split
from loguru import logger
import torchvision.transforms as transforms


def load_data(cfg):
    if cfg["val_split"] < 0 or cfg["val_split"] >= 1:
        raise ValueError(
            "val_split should be in the range of [0, 1) but got %.3f" % cfg.val_split
        )

    logger.info("Using %s dataset" % cfg["dataset"])
    if cfg["dataset"] == "cifar10":
        transform_train = transforms.Compose(
            [
                transforms.RandomCrop(32, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(
                    (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
                ),
            ]
        )

        transform_test = transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.Normalize(
                    (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
                ),
            ]
        )

        trainset = torchvision.datasets.CIFAR10(
            root="./data", train=True, download=True, transform=transform_train
        )
        trainloader = torch.utils.data.DataLoader(
            trainset, batch_size=128, shuffle=True, num_workers=2
        )

        testset = torchvision.datasets.CIFAR10(
            root="./data", train=False, download=True, transform=transform_test
        )
        testloader = torch.utils.data.DataLoader(
            testset, batch_size=100, shuffle=False, num_workers=2
        )

    return trainloader, testloader
